package com.study.classification

import org.apache.spark.ml.classification.NaiveBayes
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.{DataFrame, SparkSession}

import scala.util.Random

/**
 * 分类-朴素贝叶斯（使用鸢尾花数据集）
 *
 * @author stephen
 * @date 2019-08-27 18:40
 */
object NaiveBayesIrisDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")
    import spark.implicits._

    val path = this.getClass.getClassLoader.getResource("data/iris.data").getPath

    val data: DataFrame = spark.read
      .textFile(path)
      .map(x => {
        val random = new Random()
        val splits = x.split(",")
        val label = splits(4) match {
          case "Iris-setosa" => 0
          case "Iris-versicolor" => 1
          case "Iris-virginica" => 2
        }
        (splits(0).toDouble,
          splits(1).toDouble,
          splits(2).toDouble,
          splits(3).toDouble,
          label,
          random.nextDouble())
      }).toDF("_c0", "_c1", "_c2", "_c3", "label", "rand").sort("rand")

    val assembler = new VectorAssembler().setInputCols(Array("_c0", "_c1", "_c2", "_c3")).setOutputCol("features")
    val dataset = assembler.transform(data)
    val Array(trainingData, testData) = dataset.randomSplit(Array(0.8, 0.2))

    //bayes
    val bayes = new NaiveBayes().setFeaturesCol("features").setLabelCol("label")
    val model = bayes.fit(trainingData) //训练数据集进行训练
    val predictions = model.transform(testData)
    // 查看部分预测的结果
    predictions.show()

    val evaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")

    val accuracy = evaluator.evaluate(predictions)
    println(s"准确率： ${accuracy}")

    spark.stop()
  }
}